import tushare as ts
import datetime
import pandas as pd
from datetime import datetime, timedelta
from sqlalchemy import create_engine, text
import pymysql

pymysql.install_as_MySQLdb()
engine_ts = create_engine(
	'mysql://root:W#yne1986abc@127.0.0.1:3306/gp2024',
	pool_size=10,
	pool_recycle=1600,
	pool_pre_ping=True,
	pool_use_lifo=True,
	echo_pool=True,
	max_overflow=-1
	)

def mysql_write(df, table):
	df.to_sql(table, engine_ts, index=False, if_exists='append', chunksize=5000)
	# engine_ts.dispose()

def mysql_read(query):
	df = pd.read_sql_query(con=engine_ts.connect(), sql=text(query))
	return df
	# engine_ts.dispose()

def get_previous_workday(date):
	while date.weekday() > 4: # 0 是星期一，4 是星期五
		date -= timedelta(days=1)
	return date

def subtract_workdays(date, num_days):
	"""
	从提供的日期开始，向前推算指定的工作日数（不包括周末）。
	"""
	while num_days > 0:
		date -= timedelta(days=1)
		if date.weekday() < 5: # 如果不是周末
			num_days -= 1
	return date

def start_day(n):
	today = datetime.now().date()
	first_date = get_previous_workday(today)
	second_date = subtract_workdays(first_date, n)
	today = first_date.strftime('%Y%m%d')
	start_day = second_date.strftime('%Y%m%d')
	return today, start_day


def days(n):
	today = datetime.now().date()
	first_date = get_previous_workday(today)
	workdays = []
	current_date = first_date
	while len(workdays) < n:
	    if current_date.weekday() < 5:  # 如果不是周末
	        workdays.append(current_date)
	    current_date -= timedelta(days=1)
	# 格式化日期为yyyymmdd格式
	workdays_str = [date.strftime('%Y%m%d') for date in workdays]
	return workdays_str